import socket
import json
import os
from crypto_utils import CryptoUtils
import sys
import base64
import threading
import queue

class FTPClient:
    def __init__(self, host='127.0.0.1', port=2121):
        self.host = host
        self.port = port
        self.socket = None
        self.aes_key = None
        self.current_dir = os.getcwd()
        self.username = None
        self.chat_queue = queue.Queue()
        self.running = True

    def connect(self):
        try:
            print(f"正在连接到服务器 {self.host}:{self.port}...")
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.socket.settimeout(10)
            self.socket.connect((self.host, self.port))
            print("已连接到服务器")
            
            # 接收服务器的公钥
            print("正在接收服务器公钥...")
            server_public_key = self.socket.recv(1024)
            
            # 生成AES密钥并加密发送给服务器
            print("正在生成加密密钥...")
            self.aes_key = CryptoUtils.generate_aes_key()
            aes_key_str = base64.b64encode(self.aes_key).decode('utf-8')
            encrypted_aes_key = CryptoUtils.rsa_encrypt(server_public_key, aes_key_str)
            self.socket.send(encrypted_aes_key)
            print("密钥交换完成")
            
            return True
        except socket.timeout:
            print("连接服务器超时")
            return False
        except ConnectionRefusedError:
            print("连接被拒绝，请确保服务器正在运行")
            return False
        except Exception as e:
            print(f"连接服务器失败: {str(e)}")
            return False

    def login(self, username, password):
        try:
            self.username = username
            credentials = f"{username}:{password}"
            encrypted_credentials = CryptoUtils.aes_encrypt(self.aes_key, credentials)
            self.socket.send(encrypted_credentials)
            
            encrypted_response = self.socket.recv(1024)
            response = CryptoUtils.aes_decrypt(self.aes_key, encrypted_response)
            
            return response == "认证成功"
        except Exception as e:
            print(f"登录失败: {str(e)}")
            return False

    def start_chat_listener(self):
        """启动聊天消息监听线程"""
        def listen_for_messages():
            while True:
                try:
                    # 设置接收超时
                    self.socket.settimeout(5)  # 5秒超时
                    
                    # 接收消息长度
                    length_data = self.socket.recv(4)
                    if not length_data:
                        print("服务器连接已断开")
                        break
                        
                    message_length = int.from_bytes(length_data, byteorder='big')
                    
                    # 接收完整的加密消息
                    encrypted_message = b''
                    while len(encrypted_message) < message_length:
                        try:
                            chunk = self.socket.recv(min(4096, message_length - len(encrypted_message)))
                            if not chunk:
                                raise ConnectionError("连接已断开")
                            encrypted_message += chunk
                        except socket.timeout:
                            # 如果超时，继续尝试接收
                            continue
                        except Exception as e:
                            raise ConnectionError(f"接收消息时发生错误: {str(e)}")
                    
                    if not encrypted_message:
                        print("服务器连接已断开")
                        break
                        
                    try:
                        # 解密消息
                        message = CryptoUtils.aes_decrypt(self.aes_key, encrypted_message)
                        print(f"\n{message}")
                    except Exception as e:
                        print(f"解密消息时发生错误: {str(e)}")
                        continue
                        
                except socket.timeout:
                    # 超时后继续监听
                    continue
                except ConnectionError as e:
                    print(f"\n连接错误: {str(e)}")
                    break
                except Exception as e:
                    print(f"\n接收消息时发生错误: {str(e)}")
                    # 尝试重新连接
                    try:
                        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                        self.socket.connect((self.host, self.port))
                        print("已重新连接到服务器")
                        # 重新进行密钥交换
                        self.exchange_keys()
                    except Exception as reconnect_error:
                        print(f"重新连接失败: {str(reconnect_error)}")
                        break
                        
        # 启动监听线程
        listener_thread = threading.Thread(target=listen_for_messages)
        listener_thread.daemon = True
        listener_thread.start()

    def send_chat_message(self, message):
        """发送聊天消息"""
        try:
            command = f"chat:{message}"
            encrypted_command = CryptoUtils.aes_encrypt(self.aes_key, command)
            
            # 发送消息长度
            message_length = len(encrypted_command)
            length_data = message_length.to_bytes(4, byteorder='big')
            self.socket.send(length_data)
            
            # 发送加密消息
            self.socket.send(encrypted_command)
            print(f"你: {message}")  # 显示自己发送的消息
        except Exception as e:
            print(f"发送聊天消息失败: {str(e)}")

    def send_command(self, command):
        try:
            encrypted_command = CryptoUtils.aes_encrypt(self.aes_key, command)
            
            # 发送消息长度
            message_length = len(encrypted_command)
            length_data = message_length.to_bytes(4, byteorder='big')
            self.socket.send(length_data)
            
            # 发送加密消息
            self.socket.send(encrypted_command)
            
            # 接收响应长度
            length_data = self.socket.recv(4)
            if not length_data:
                return "连接已断开"
                
            response_length = int.from_bytes(length_data, byteorder='big')
            
            # 接收完整的加密响应
            encrypted_response = b''
            while len(encrypted_response) < response_length:
                chunk = self.socket.recv(min(4096, response_length - len(encrypted_response)))
                if not chunk:
                    break
                encrypted_response += chunk
                
            if encrypted_response:
                response = CryptoUtils.aes_decrypt(self.aes_key, encrypted_response)
                return response
            return "未收到响应"
            
        except Exception as e:
            return f"发送命令失败: {str(e)}"

    def list_directory(self):
        response = self.send_command('list')
        try:
            files = json.loads(response)
            return files
        except:
            return response

    def change_directory(self, directory):
        return self.send_command(f'cd {directory}')

    def change_to_parent_directory(self):
        return self.send_command('cd..')

    def get_current_directory(self):
        return self.send_command('pwd')

    def create_directory(self, directory):
        return self.send_command(f'mkdir {directory}')

    def get_file(self, filename):
        response = self.send_command(f'get {filename}')
        if response.startswith("文件内容:"):
            content = response[6:]
            with open(filename, 'wb') as f:
                f.write(content.encode())
            return f"文件 {filename} 下载成功"
        return response

    def put_file(self, filename):
        try:
            if not os.path.exists(filename):
                return "文件不存在"
            
            with open(filename, 'rb') as f:
                content = f.read()
            
            response = self.send_command(f'put {filename}')
            return response
        except Exception as e:
            return f"上传文件失败: {str(e)}"

    def get_file_attributes(self, filename):
        response = self.send_command(f'attrib {filename}')
        try:
            return json.loads(response)
        except:
            return response

    def rename_file(self, old_name, new_name):
        return self.send_command(f'rename {old_name} {new_name}')

    def create_file(self, filename):
        return self.send_command(f'create {filename}')

    def add_user(self, username, password):
        return self.send_command(f'add {username} {password}')

    def close(self):
        self.running = False
        if self.socket:
            self.socket.close()

def main():
    client = FTPClient()
    if not client.connect():
        return

    print("欢迎使用FTP客户端")
    username = input("请输入用户名: ")
    password = input("请输入密码: ")

    if not client.login(username, password):
        print("登录失败")
        client.close()
        return

    print("登录成功！")
    print("可用命令：")
    print("list - 列出当前目录文件")
    print("cd <目录名> - 进入指定目录")
    print("cd.. - 返回上级目录")
    print("pwd - 显示当前目录")
    print("mkdir <目录名> - 创建目录")
    print("get <文件名> - 下载文件")
    print("put <文件名> - 上传文件")
    print("attrib <文件名> - 查看文件属性")
    print("rename <旧文件名> <新文件名> - 重命名文件")
    print("create <文件名> - 创建文件")
    print("add <用户名> <密码> - 添加用户")
    print("chat <消息> - 发送聊天消息")
    print("exit - 退出程序")

    # 启动聊天消息监听
    client.start_chat_listener()

    while True:
        try:
            command = input("\n请输入命令: ")
            if command.lower() == 'exit':
                break

            cmd_parts = command.split()
            cmd = cmd_parts[0].lower()

            if cmd == 'chat':
                if len(cmd_parts) > 1:
                    message = ' '.join(cmd_parts[1:])
                    client.send_chat_message(message)
                else:
                    print("用法: chat <消息>")
            elif cmd == 'list':
                result = client.list_directory()
                if isinstance(result, list):
                    for item in result:
                        print(item)
                else:
                    print(result)
            elif cmd == 'cd':
                if len(cmd_parts) > 1:
                    print(client.change_directory(cmd_parts[1]))
                else:
                    print("用法: cd <目录名>")
            elif cmd == 'cd..':
                print(client.change_to_parent_directory())
            elif cmd == 'pwd':
                print(client.get_current_directory())
            elif cmd == 'mkdir':
                if len(cmd_parts) > 1:
                    print(client.create_directory(cmd_parts[1]))
                else:
                    print("用法: mkdir <目录名>")
            elif cmd == 'get':
                if len(cmd_parts) > 1:
                    print(client.get_file(cmd_parts[1]))
                else:
                    print("用法: get <文件名>")
            elif cmd == 'put':
                if len(cmd_parts) > 1:
                    print(client.put_file(cmd_parts[1]))
                else:
                    print("用法: put <文件名>")
            elif cmd == 'attrib':
                if len(cmd_parts) > 1:
                    print(client.get_file_attributes(cmd_parts[1]))
                else:
                    print("用法: attrib <文件名>")
            elif cmd == 'rename':
                if len(cmd_parts) > 2:
                    print(client.rename_file(cmd_parts[1], cmd_parts[2]))
                else:
                    print("用法: rename <旧文件名> <新文件名>")
            elif cmd == 'create':
                if len(cmd_parts) > 1:
                    print(client.create_file(cmd_parts[1]))
                else:
                    print("用法: create <文件名>")
            elif cmd == 'add':
                if len(cmd_parts) > 2:
                    print(client.add_user(cmd_parts[1], cmd_parts[2]))
                else:
                    print("用法: add <用户名> <密码>")
            else:
                print("未知命令")

        except Exception as e:
            print(f"执行命令时发生错误: {str(e)}")

    client.close()

if __name__ == '__main__':
    main() 